{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Issue [#118](https://github.com/bsorrentino/langgraph4j/issues/118) by [quovadis1234](https://github.com/quovadis1234)\n",
    "\n",
    "Verify \"**How to interrupt graph while streaming node output?**\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "var userHomeDir = System.getProperty(\"user.home\");\n",
    "var localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "var langchain4jVersion = \"1.0.1\";\n",
    "var langchain4jbeta = \"1.0.1-beta6\";\n",
    "var langgraph4jVersion = \"1.6-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove installed package from Jupiter cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos\n",
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-ollama:\\{langchain4jbeta}\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize Logger**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    java.util.logging.LogManager.getLogManager().readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"issue118\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use StreamingChatGenerator in Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Graph State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import org.bsc.langgraph4j.prebuilt.MessagesState;\n",
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "\n",
    "class State extends MessagesState<ChatMessage> {\n",
    "    public State( Map<String, Object> initData ) {\n",
    "        super( initData );\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.model.chat.request.ChatRequest;\n",
    "import dev.langchain4j.model.ollama.OllamaStreamingChatModel;\n",
    "import org.bsc.langgraph4j.action.NodeAction;\n",
    "import org.bsc.langgraph4j.langchain4j.generators.StreamingChatGenerator;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.LC4jStateSerializer;\n",
    "import org.bsc.langgraph4j.streaming.StreamingOutput;\n",
    "import org.bsc.langgraph4j.StateGraph;\n",
    "import java.util.Map;\n",
    "\n",
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;\n",
    "\n",
    "\n",
    "var model = OllamaStreamingChatModel.builder()\n",
    "        .baseUrl(\"http://localhost:11434\")\n",
    "        .temperature(0.0)\n",
    "        .logRequests(true)\n",
    "        .logResponses(true)\n",
    "        .modelName(\"llama3.1:latest\")\n",
    "        .build();\n",
    "\n",
    "NodeAction<State> calculationNode = state -> {\n",
    "\n",
    "        log.trace(\"calculationNode: {}\", state.messages());\n",
    "\n",
    "        var generator = StreamingChatGenerator.<State>builder()\n",
    "        .mapResult(response -> Map.of(\"messages\", response.aiMessage()))\n",
    "        .startingNode(\"calculation\")\n",
    "        .startingState(state)\n",
    "        .build();\n",
    "\n",
    "        var request = ChatRequest.builder()\n",
    "        .messages(state.messages())\n",
    "        .build();\n",
    "\n",
    "        model.chat(request, generator.handler());\n",
    "\n",
    "        return Map.of(\"_streaming_messages\", generator);\n",
    "};\n",
    "\n",
    "NodeAction<State> summaryNode = state -> {\n",
    "        log.trace(\"summaryNode: {}\", state.messages());\n",
    "\n",
    "        var lastMessage = state.lastMessage().orElseThrow();\n",
    "\n",
    "        return Map.of();\n",
    "};\n",
    "\n",
    "var stateSerializer = new LC4jStateSerializer<State>(State::new);\n",
    "\n",
    "// Define Graph\n",
    "var workflow = new StateGraph<State>(State.SCHEMA, stateSerializer)\n",
    "        .addNode(\"calculation\", node_async(calculationNode))\n",
    "        .addNode(\"summary\", node_async(summaryNode))\n",
    "        .addEdge(START, \"calculation\")\n",
    "        .addEdge(\"calculation\", \"summary\" )\n",
    "        .addEdge(\"summary\", END);\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "__START__ - UserMessage { name = null contents = [TextContent { text = \"generate a UUID for me\" }] } \n",
      "StreamingOutput{node=calculation, chunk=` } \n",
      "StreamingOutput{node=calculation, chunk=4 } \n",
      "StreamingOutput{node=calculation, chunk=c } \n",
      "StreamingOutput{node=calculation, chunk=9 } \n",
      "StreamingOutput{node=calculation, chunk=f } \n",
      "StreamingOutput{node=calculation, chunk=6 } \n",
      "StreamingOutput{node=calculation, chunk=a } \n",
      "StreamingOutput{node=calculation, chunk=1 } \n",
      "StreamingOutput{node=calculation, chunk=d } \n",
      "StreamingOutput{node=calculation, chunk=- } \n",
      "StreamingOutput{node=calculation, chunk=5 } \n",
      "StreamingOutput{node=calculation, chunk=b } \n",
      "StreamingOutput{node=calculation, chunk=3 } \n",
      "StreamingOutput{node=calculation, chunk=e } \n",
      "StreamingOutput{node=calculation, chunk=- } \n",
      "StreamingOutput{node=calculation, chunk=4 } \n",
      "StreamingOutput{node=calculation, chunk=c } \n",
      "StreamingOutput{node=calculation, chunk=8 } \n",
      "StreamingOutput{node=calculation, chunk=f } \n",
      "StreamingOutput{node=calculation, chunk=-b } \n",
      "StreamingOutput{node=calculation, chunk=2 } \n",
      "StreamingOutput{node=calculation, chunk=e } \n",
      "StreamingOutput{node=calculation, chunk=7 } \n",
      "StreamingOutput{node=calculation, chunk=- } \n",
      "StreamingOutput{node=calculation, chunk=0 } \n",
      "StreamingOutput{node=calculation, chunk=f } \n",
      "StreamingOutput{node=calculation, chunk=1 } \n",
      "StreamingOutput{node=calculation, chunk=c } \n",
      "StreamingOutput{node=calculation, chunk=5 } \n",
      "StreamingOutput{node=calculation, chunk=e } \n",
      "StreamingOutput{node=calculation, chunk=8 } \n",
      "StreamingOutput{node=calculation, chunk=a } \n",
      "StreamingOutput{node=calculation, chunk=8 } \n",
      "StreamingOutput{node=calculation, chunk=5 } \n",
      "StreamingOutput{node=calculation, chunk=eb } \n",
      "StreamingOutput{node=calculation, chunk=` } \n",
      "calculation - AiMessage { text = \"`4c9f6a1d-5b3e-4c8f-b2e7-0f1c5e8a85eb`\" toolExecutionRequests = [] } \n",
      "summary - AiMessage { text = \"`4c9f6a1d-5b3e-4c8f-b2e7-0f1c5e8a85eb`\" toolExecutionRequests = [] } \n",
      "__END__ - AiMessage { text = \"`4c9f6a1d-5b3e-4c8f-b2e7-0f1c5e8a85eb`\" toolExecutionRequests = [] } \n"
     ]
    }
   ],
   "source": [
    "\n",
    "var app = workflow.compile();\n",
    "\n",
    "for( var out : app.stream( Map.of( \"messages\", UserMessage.from( \"generate a UUID for me\")) ) ) {\n",
    "        \n",
    "        if( out instanceof StreamingOutput streaming ) {\n",
    "                log.info( \"StreamingOutput{node={}, chunk={} }\", streaming.node(), streaming.chunk() );\n",
    "        }\n",
    "        else {\n",
    "                log.info( \"{} - {}\", out.node(), out.state().lastMessage().orElseThrow() );\n",
    "        }\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "22.0.2+9-70"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
